#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
PostgreSQL复制模块
提供复制管理功能：配置主从复制、查看复制状态等
"""

import os
import re
from typing import Optional

from ...core.utils import logger, run_command, get_input_with_default
from ...core import config

class PostgresqlReplication:
    """PostgreSQL复制管理类"""
    
    def change_replica_master_ip(self):
        """修改从库指向的主库IP"""
        if not self.config_path:
            logger.error("找不到PostgreSQL配置路径")
            return
            
        new_master_ip = get_input_with_default("输入新主控IP", "")
        repl_password = get_input_with_default(f"输入复制用户 '{config.DEFAULT_REPLICATION_USER}' 的密码", "")
        
        if not new_master_ip:
            logger.error("IP地址无效")
            return

        if not repl_password:
            logger.warning("未提供密码，将尝试使用现有配置的密码")
            
        postgresql_conf = os.path.join(self.config_path, "postgresql.conf")
        postgresql_auto_conf = os.path.join(self.config_path, "postgresql.auto.conf")
        
        # 修改配置文件
        self.update_primary_conninfo(postgresql_conf, new_master_ip, repl_password)
        self.update_primary_conninfo(postgresql_auto_conf, new_master_ip, repl_password)
        
        # 重启PostgreSQL服务
        logger.info("正在重启PostgreSQL以应用新的配置...")
        self.restart()
    
    def update_primary_conninfo(self, config_file: str, new_ip: str, password: Optional[str] = None):
        """修改配置文件中的primary_conninfo参数
        
        Args:
            config_file: 配置文件路径
            new_ip: 新的主服务器IP
            password: 复制用户密码，如果为None则尝试保留现有配置的密码
        """
        if not os.path.exists(config_file):
            logger.error(f"配置文件 {config_file} 不存在")
            return
            
        try:
            # 读取文件内容
            with open(config_file, "r") as file:
                lines = file.readlines()
            
            # 检查是否存在primary_conninfo行，并提取现有密码
            existing_password = None
            if not password:
                for line in lines:
                    if "primary_conninfo" in line and "password=" in line:
                        # 尝试提取现有密码
                        password_match = re.search(r"password=([^ '\"]+)", line)
                        if password_match:
                            existing_password = password_match.group(1)
                            logger.info("已从现有配置提取密码")
                            break
            
            # 如果没有提供新密码且未找到现有密码，提示错误
            if not password and not existing_password:
                logger.error("无法确定复制用户密码，请提供密码")
                return
                
            # 使用找到的密码或提供的新密码
            final_password = password or existing_password
            
            # 创建新的primary_conninfo行
            new_line = (f"primary_conninfo = 'host={new_ip} "
                       f"user={config.DEFAULT_REPLICATION_USER} "
                       f"password={final_password} "
                       f"application_name={config.DEFAULT_APPLICATION_NAME}'\n")
                
            # 标记是否找到并替换了行
            found = False
            with open(config_file, "w") as file:
                for line in lines:
                    if "primary_conninfo" in line:
                        file.write(new_line)
                        found = True
                    else:
                        file.write(line)
                        
            # 如果没有找到现有配置，追加新行
            if not found:
                with open(config_file, "a") as file:
                    file.write(new_line)
                    
            logger.info(f"已更新 {config_file} 中的 primary_conninfo")
        except Exception as e:
            logger.error(f"更新配置文件 {config_file} 时出错: {e}")
    
    def show_master_replication_status(self):
        """查看主库流复制状态"""
        logger.info("主库复制状态:")
        
        try:
            run_command([
                "sudo", "-u", "postgres", "psql", 
                "-c", "\\x", 
                "-c", "SELECT * FROM pg_stat_replication;"
            ])
        except Exception as e:
            logger.error(f"检索主库复制状态时出错: {e}")
    
    def show_slave_replication_status(self):
        """查看从库流复制状态"""
        logger.info("从库复制状态:")
        
        try:
            run_command([
                "sudo", "-u", "postgres", "psql", 
                "-c", "\\x", 
                "-c", "SELECT * FROM pg_stat_wal_receiver;"
            ])
        except Exception as e:
            logger.error(f"检索从库复制状态时出错: {e}")
    
    def add_ip_to_pghba(self):
        """一键添加pghba.conf信任IP"""
        ip_addresses = get_input_with_default("请输入需要添加的IP地址（多个地址用逗号分隔）", "").split(',')
        
        # 清理IP地址的空白字符
        ip_addresses = [ip.strip() for ip in ip_addresses if ip.strip()]
        
        if not ip_addresses:
            logger.error("未提供有效的IP地址")
            return
            
        # pg_hba.conf 文件路径
        pghba_path = f'/etc/postgresql/{self.pg_version}/main/pg_hba.conf'
        
        # 定义要添加的规则模板
        rules_to_add = []
        for ip in ip_addresses:
            rules_to_add.extend([
                f"host    replication     all             {ip}/32       trust\n",
                f"host    all             all             {ip}/32       trust\n",
                f"host    replication     replicator      {ip}/32       trust\n",
                f"host    replication     postgres        {ip}/32       trust\n",
                f"host    replication     repmgr          {ip}/32       trust\n"
            ])
            
        try:
            # 打开 pg_hba.conf 文件并追加规则
            with open(pghba_path, 'a') as pghba_file:
                pghba_file.writelines(rules_to_add)
            logger.info("IP地址成功添加到 pg_hba.conf 文件")
            
            # 重新加载PostgreSQL配置
            self.reload()
            logger.info("已重载数据库配置文件")
        except Exception as e:
            logger.error(f"修改 pg_hba.conf 文件时出错: {e}") 